# coding = utf-8

'''
验证读取数据是否有问题
'''

import sys

import click
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib2 import Path
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from tqdm import tqdm

import utils.checkpoint as cp
from dataset import KiTS19
from dataset.transform import MedicalTransform
from loss import GeneralizedDiceLoss
from loss.util import class2one_hot
from network.m_net import UNet
from utils.metrics import Evaluator
from utils.vis import imshow
import matplotlib.pyplot as plt
import os
import cv2

IMAGE_PATH = "/datasets/DongbeiDaxue/chengkunv2"

def validate_read():
    dataset = KiTS19("/datasets/DongbeiDaxue/chengkunv2", stack_num=3, spec_classes=[0, 1, 2], img_size=(512, 512),
                     use_roi=False, roi_file=None, roi_error_range=5,
                     train_transform=None, valid_transform=None)

    subset = dataset.valid_dataset
    sampler = SequentialSampler(subset)
    data_loader = DataLoader(subset, batch_size=1, sampler=sampler,
                             num_workers=1, pin_memory=True)

    for batch_idx, data in enumerate(data_loader):
        case_path = os.path.join(IMAGE_PATH, "case_00035/imaging")
        imgs, labels, idx = data['image'], data['label'], data['index']
        imgs = imgs.squeeze().numpy()
        plt.subplot(1, 4, 2)
        plt.imshow(imgs[0], cmap="gray")
        plt.subplot(1, 4, 3)
        plt.imshow(imgs[1], cmap="gray")
        plt.subplot(1, 4, 4)
        plt.imshow(imgs[2], cmap="gray")
        file_name = os.path.join(case_path, "{}.npy".format(str(batch_idx).zfill(3)))
        origion = np.load(file_name)
        plt.subplot(1, 4, 1)
        plt.imshow(origion, cmap="gray")
        plt.show()



if __name__ == '__main__':
    validate_read()